package onion;

import java.io.*;
import java.net.*;
import java.util.*;
import java.security.*;
import java.security.spec.*;
import javax.crypto.*;
import javax.crypto.spec.*;
import onion.*;
/**
 *
 * @author xDEMONx
 * Узел в сети
 */
public class Node {
    //IP-адрес сервера каталогов; Имя конфигурационного файла, формата:
    //     -IP-адрес сервера каталогов
    //     -Его порт
    //     -Желаемый порт для узла
    static String IPSereverDB = "127.0.0.1", ConfFile = "node.conf";    
    //Порт сервера каталогов; Порт узла
    static int PortServerDB = 1315, myPort = 1617;        
    KeyPair keyPair;    
    //Поток обрабатывающий проходящий трафик
    class WorkerNode extends Thread{
        Socket s, nextSocket;
        String IPnext;
        int PortNext;
        SecretKey sk;
        byte[] buf;
        boolean nextOpen = false;
        public WorkerNode(Socket soc) {
            s = soc;
        }
        public void run(){
            buf = new byte[OnionConst.max_size];
            try{
                //Построение пути и получение ключа для симметричного шифра
                DataInputStream in = new DataInputStream(s.getInputStream());
                DataOutputStream out = null;
                int count = in.read(buf);                
                //Байты для дешифровки
                byte[] CT = Arrays.copyOfRange(buf, 0, count);                                
                int rounds =  (count+OnionConst.block_size_decrypt-1)/OnionConst.block_size_decrypt;
                int lenPT = 0;
                byte[][] ListPT = new byte[rounds][];
                //Дешифруем закрытым ключом
                Cipher RSA_D = Cipher.getInstance(OnionConst.PublicKeyEncrypt); 
                RSA_D.init(Cipher.DECRYPT_MODE, keyPair.getPrivate());                                    
                int pos_CT = 0;
                for(int i = 0; i < rounds; i++){                    
                    ListPT[i] = RSA_D.doFinal(Arrays.copyOfRange(CT, pos_CT, pos_CT+OnionConst.block_size_decrypt));
                    lenPT+=ListPT[i].length;
                    pos_CT+=OnionConst.block_size_decrypt;                    
                }                
                byte[] PT  = new byte[lenPT];
                for(int i=0,cur = 0; i < rounds; i++){
                    for(int j = 0; j < ListPT[i].length; j++, cur++){
                        PT[cur] = ListPT[i][j];
                    }
                }                
                //Является ли данный узел последним в цепочке
                if(PT.length > OnionConst.SKE_keysize/8){
                    //Есть следующий. Получаем его адрес
                    nextOpen = true;
                    byte[] IP = {PT[0],PT[1],PT[2],PT[3]};
                    IPnext = Inet4Address.getByAddress(IP).getHostAddress();
                    PortNext = ((PT[5] & 0xff)<<8) | (PT[4]&0xff);                    
                    //Извлекаем симметричный ключ
                    byte[] SK_bytes = new byte[OnionConst.SKE_keysize/8];
                    for(int i = 0; i < SK_bytes.length; i++){
                        SK_bytes[i] = PT[6+i];
                    }
                    sk = new SecretKeySpec(SK_bytes,OnionConst.SymmeetricKeyEncrypt);                    
                    //Передаём пакет дальше
                    nextSocket = new Socket(IPnext, PortNext);
                    out = new DataOutputStream(nextSocket.getOutputStream());
                    out.write(PT, 6 + OnionConst.SKE_keysize/8, 
                            PT.length - 6 - OnionConst.SKE_keysize/8);
                    out.flush();                    
                }
                else{
                    //Это последний узел. Извлекаем симметричный ключ
                    nextOpen = false;
                    sk = new SecretKeySpec(PT,OnionConst.SymmeetricKeyEncrypt);
                }             
                //Создаём дешифратор для симметричного алгоритма
                Cipher AES_D = Cipher.getInstance(OnionConst.SymmeetricKeyEncrypt + 
                        OnionConst.SymmetricEncryptMode);                
                boolean good = true;   //Пока всё хорошо выполняется цикл
                while(good){
                    //Читаем очередную порцию байт
                    count = in.read(buf);
                    //Байты для дешифровки
                    byte IV[] = new byte[OnionConst.SKE_block_size_byte];
                    CT = new byte[count - IV.length];                
                    for(int i = 0; i < count; i++){
                        if(i < IV.length){
                            IV[i] = buf[i];
                        }else{
                            CT[i - IV.length] = buf[i];
                        }
                    }
                    //Расшифровываем симметричным алгоритмом
                    AES_D.init(Cipher.DECRYPT_MODE, sk, new IvParameterSpec(IV));
                    PT = AES_D.doFinal(CT);
                    //Получаем адрес следующего узла
                    byte[] IP = {PT[0],PT[1],PT[2],PT[3]};
                    String tmpIP = Inet4Address.getByAddress(IP).getHostAddress();
                    int tmpPort = ((PT[5] & 0xff)<<8) | (PT[4]&0xff);
                    if(nextOpen==false){
                        nextSocket = new Socket(tmpIP, tmpPort);
                        out = new DataOutputStream(nextSocket.getOutputStream());                        
                    }                    
                    if(nextOpen == false || (IPnext.equals(tmpIP) && PortNext==tmpPort)){
                        //Шлём дальше
                        out.write(PT, 6, PT.length - 6);
                        out.flush();
                        if(nextOpen == false){
                            out.close();
                            nextSocket.close();
                        }
                    }
                    else{
                        System.out.printf("Error: Different destination address."
                                + "I get \"%s:%d\", but I have \"%s:%d\"\n",
                                tmpIP,tmpPort,IPnext,PortNext);
                        good = false;
                    }                        
                }                
            }
            catch(Exception e){
                System.out.println("Error : " + e);
            }                
        }
    }
    public static void main(String[] args) throws IOException, NoSuchAlgorithmException{
        new Node().run();
    }
    void run() throws IOException, NoSuchAlgorithmException{
        try{
            //Попытка считать конфигурацию из файла
            Scanner fileIn = new Scanner(new FileReader(ConfFile));
            IPSereverDB = fileIn.nextLine();
            PortServerDB = fileIn.nextInt();
            myPort = fileIn.nextInt();
            fileIn.close();
        }
        catch(FileNotFoundException e){    
            //Файла не нашлось, пишем сами
            PrintWriter fileOut = new PrintWriter(new FileWriter(ConfFile));
            fileOut.println(IPSereverDB);            
            fileOut.println(PortServerDB);
            fileOut.println(myPort);
            fileOut.close();
        }
        //Ищем свободный порт и создаём сокет
        boolean socket_created = false;
        ServerSocket nodeSoket = null;
        while(!socket_created){
            try{
                nodeSoket = new ServerSocket(myPort);
                System.out.println("OnionNode is started on port "+myPort);
                socket_created = true;
            }
            catch(Exception e){
                myPort++;
            }
        }        
        KeyPairGenerator kpg = KeyPairGenerator.getInstance(OnionConst.PublicKeyEncrypt);
        kpg.initialize(OnionConst.PKE_keysize);
        keyPair = kpg.genKeyPair();                    
        boolean isReady = false;
        //Сообщаем информацию об узле серверу каталогов
        try{      
           Socket servDB = new Socket(IPSereverDB, PortServerDB);
           ObjectOutputStream out = new ObjectOutputStream(servDB.getOutputStream());                                 
           out.writeUTF("add");
           NodeAddress na = new NodeAddress();
           na.IP = servDB.getLocalAddress().getHostAddress();
           na.port = myPort;
           na.PK = keyPair.getPublic();                       
           out.writeObject(na);
           out.close();
           servDB.close();
           isReady = true;
        }
        catch(Exception e){
            System.out.println("init error: "+e);
        }        
        if(isReady){
            System.out.println("Node ready to work");
            while(true){
                Socket s = nodeSoket.accept();
                WorkerNode wn = new WorkerNode(s);
                wn.setDaemon(true);
                wn.start();
            }
        }
        else{
            System.out.println("Error: Can not connect to serverDB");
        }            
    }
}
